#!/usr/bin/env python3

# -------- navigation_simulate.py ---------

import rospy
import math
import tf
import zlib  # 用于计算校验和
from geometry_msgs.msg import TransformStamped, Twist
from nav_msgs.msg import OccupancyGrid
from tf.transformations import quaternion_from_euler
from geometry_msgs.msg import PoseWithCovarianceStamped
from tf.transformations import euler_from_quaternion


class NavSimulate:
    def __init__(self):
        rospy.init_node("nav_simulate", anonymous=False)

        # 初始化位姿
        self.x = rospy.get_param("~initial_x", 0.0)
        self.y = rospy.get_param("~initial_y", 0.0)
        self.yaw = rospy.get_param("~initial_yaw", 0.0)

        # 存储速度
        self.linear_velocity = 0.0
        self.angular_velocity = 0.0

        # 全局地图
        self.global_map = None
        self.map_feature = None  # 地图特征值

        # 上次时间戳
        self.last_time = rospy.Time.now()

        rospy.Subscriber("/cmd_vel", Twist, self.cmd_vel_callback)
        rospy.Subscriber("/map", OccupancyGrid, self.map_callback)
        self.local_map_pub = rospy.Publisher("/local_map", OccupancyGrid, queue_size=10)

        rospy.Subscriber(
            "/initialpose", PoseWithCovarianceStamped, self.pose_reset_callback
        )

        self.br = tf.TransformBroadcaster()

        self.timer = rospy.Timer(rospy.Duration(1.0 / 20.0), self.timer_callback)

    def cmd_vel_callback(self, msg):
        self.linear_velocity = msg.linear.x
        self.angular_velocity = msg.angular.z

    def map_callback(self, msg):
        # 计算新地图的特征值（使用CRC32校验和）
        new_feature = zlib.crc32(bytes(msg.data))
        
        # 只有当特征值变化时才更新地图
        if self.map_feature != new_feature:
            self.global_map = msg
            self.map_feature = new_feature
            rospy.logdebug(f"Global map updated with new feature: {new_feature}")
        else:
            rospy.logdebug("Received duplicate map, skipping update")

    def timer_callback(self, event):
        current_time = rospy.Time.now()
        dt = (current_time - self.last_time).to_sec()
        if dt > 0:
            # 更新位姿
            self.x += self.linear_velocity * math.cos(self.yaw) * dt
            self.y += self.linear_velocity * math.sin(self.yaw) * dt
            self.yaw += self.angular_velocity * dt

        self.last_time = current_time

        # 发布 TF 变换
        self.publish_transform()

        # 发布局部地图
        if self.global_map:
            local_map = self.create_local_map()
            self.local_map_pub.publish(local_map)

    def publish_transform(self):
        # 发布 map 到 body 的变换
        q = quaternion_from_euler(0, 0, self.yaw)
        self.br.sendTransform((self.x, self.y, 0), q, rospy.Time.now(), "body", "map")

    def create_local_map(self):
        local_map = OccupancyGrid()
        info = self.global_map.info
        resolution = info.resolution

        # 设置 local_map 的基本参数
        width = max(1, int(12.0 / resolution))
        height = max(1, int(12.0 / resolution))

        local_map.header.stamp = rospy.Time.now()
        local_map.header.frame_id = "body"
        local_map.info.resolution = resolution
        local_map.info.width = width
        local_map.info.height = height
        local_map.info.origin.position.x = -4.0
        local_map.info.origin.position.y = -4.0
        local_map.info.origin.orientation.w = 1.0  # identity rotation

        data = []

        for row in range(height):
            for col in range(width):
                # body 坐标系下的坐标
                x_prime = -4.0 + (col + 0.5) * resolution
                y_prime = -4.0 + (row + 0.5) * resolution

                # 转换到 map 坐标系
                cos_yaw = math.cos(self.yaw)
                sin_yaw = math.sin(self.yaw)
                x_map = self.x + x_prime * cos_yaw - y_prime * sin_yaw
                y_map = self.y + x_prime * sin_yaw + y_prime * cos_yaw

                # 转换到全局地图的 cell 坐标
                origin_x = info.origin.position.x
                origin_y = info.origin.position.y
                map_width = info.width
                map_height = info.height

                cell_x = (x_map - origin_x) / resolution
                cell_y = (y_map - origin_y) / resolution

                cell_col = int(math.floor(cell_x))
                cell_row = int(math.floor(cell_y))

                # 检查是否在地图范围内
                if 0 <= cell_col < map_width and 0 <= cell_row < map_height:
                    index = cell_row * map_width + cell_col
                    value = self.global_map.data[index]
                else:
                    value = -1  # 未知

                data.append(value)

        local_map.data = data
        return local_map

    def pose_reset_callback(self, msg):
        # 提取位置信息
        self.x = msg.pose.pose.position.x
        self.y = msg.pose.pose.position.y

        # 提取姿态信息并转换为欧拉角
        orientation_q = msg.pose.pose.orientation
        orientation_list = [
            orientation_q.x,
            orientation_q.y,
            orientation_q.z,
            orientation_q.w,
        ]
        _, _, self.yaw = euler_from_quaternion(orientation_list)

        # 更新时间戳以避免 dt 异常
        self.last_time = rospy.Time.now()


if __name__ == "__main__":
    try:
        ns = NavSimulate()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass